from functools import cmp_to_key

import pandas as pd

from IndexCollector import IndexCollector
from libs.Objects import Objects
from StockDownloader import StockDownloader
from StockToFile import StockToFile
from libs.StockBase import StockBase

DAY_FORMAT = '%Y-%m-%d'

stockBase = StockBase()


class DayResult(object):
    def __init__(self, tradingDate):
        self.day = tradingDate
        self.green_count = None
        self.red_count = None
        self.limit_up = []
        self.limit_down = []
        self.break_down = []
        self.shsz_amount = 0

    def getDay(self):
        return self.day

    def get_red_count(self):
        return self.red_count

    def set_red_count(self, value):
        self.red_count = value

    def get_green_count(self):
        return self.green_count

    def set_green_count(self, value):
        self.green_count = value

    def get_limit_down(self):
        return self.limit_down

    def add_limit_down(self, value):
        self.limit_down.append(value)

    def get_break_down(self):
        return self.break_down

    def add_break_down(self, value):
        self.break_down.append(value)

    def get_limit_up(self):
        return self.limit_up

    def add_limit_up(self, value):
        self.limit_up.append(value)

    def setDayLimitUp(self, df_value):
        self.day_limit_up = df_value

    # date df code,strong
    def getDayLimitUp(self):
        return self.day_limit_up

    def setShSzAmount(self, value):
        self.shsz_amount = value

    def getShSzAmount(self):
        return self.shsz_amount

    def toString(self):
        msg = '{day} red:{red},green:{green},limitUp:{limitUp},limitDown:{limitDown},breakDown:{breakDown},amount:{amount}'
        print(msg.format(day=self.day,
                         red=self.red_count,
                         breakDown=self.break_down,
                         green=self.green_count,
                         limitUp=len(self.limit_up),
                         limitDown=len(self.limit_down),
                         amount=self.shsz_amount))


class StockProcessor(object):
    CALC_PRE_DAYS = 20

    def __init__(self):
        pass

    def process(self, startDay, endDay):
        self.caclStartDay = Objects.getDiffDay(startDay, -1 * StockProcessor.CALC_PRE_DAYS)
        self.calcTradingDays = StockBase().getTradingDays(self.caclStartDay, endDay)

        df = self.preProcess(startDay, endDay)
        all_result = self.doProcess(df, startDay, endDay)
        for result in all_result:
            self.postProcess(result)

    def postProcess(self, result):
        df = result.getDayLimitUp()

        breakCodes = "，".join([code for code in result.get_break_down()])

        temp1Df = df[df['strong'] == 'd1u1']
        d1u1 = "，".join([code for code in temp1Df['code']])

        temp2Df = df[(df['strong'] >= 'd2u2') & (df['strong'] <= 'd3u2')]
        d2u2 = "，".join([row.code + row.strong for row in temp2Df.itertuples()])

        temp3Df = df[(df['strong'] >= 'd3u3') & (df['strong'] <= 'd4u3')]
        d3u3 = "，".join([row.code + row.strong for row in temp3Df.itertuples()])

        temp4Df = df[df['strong'] >= 'd4u4']
        gtD3u3 = "，".join([row.code + row.strong for row in temp4Df.itertuples()])

        stockToFile = StockToFile()
        # stockToFile.appendFile(result.getDay(), data)
        stockToFile.to_db(
            [Objects.to_db_date(result.getDay()), result.get_red_count(), result.get_green_count(),
             len(result.get_limit_up()), len(result.get_limit_down()), breakCodes, d1u1, d2u2, d3u3, gtD3u3,
             result.getShSzAmount()])

    def preProcess(self, startDay, endDay):
        files = Objects.getAllFileByDay('data/trading', self.caclStartDay, endDay)
        df_arry = [pd.read_csv(file) for file in files]
        all_df = pd.concat(df_arry, axis=0, ignore_index=True)
        print("read all csv")
        ## process
        all_df = all_df.drop(all_df[all_df['isST'] == 1].index)
        all_df['limitUp'] = all_df.apply(self.isLimitUp, axis=1)
        all_df['limitDown'] = all_df.apply(self.isLimitDown, axis=1)
        all_df['breakDown'] = all_df.apply(self.isBreakDown, axis=1)
        ## test
        test = False
        if test:
            all_df = all_df[all_df['code'] == 'sh.600128']
            print(all_df.sort_values(by='date', ascending=False))
        return all_df

    def doProcess(self, all_df, startDay, endDay):
        tradingDays = [d for d in list(filter(lambda day: startDay <= day and day <= endDay, self.calcTradingDays))]

        all_result = []
        for day in tradingDays:
            if not StockToFile().if_exist(day):
                result = self.doProcessByDay(all_df, day)
                all_result.append(result)

        return all_result

    def doProcessByDay(self, all_df, tradingDay):
        print("doProcessByDay:" + str(tradingDay))
        dayResult = DayResult(tradingDay)
        dayResult.set_red_count(all_df[(all_df['date'] == tradingDay) & (all_df['pctChg'] > 0)].shape[0])
        dayResult.set_green_count(all_df[(all_df['date'] == tradingDay) & (all_df['pctChg'] < 0)].shape[0])

        codes = all_df[(all_df['date'] == tradingDay) & all_df['breakDown']]['code'].tolist()
        [dayResult.add_break_down(c) for c in codes]

        codes = all_df[(all_df['date'] == tradingDay) & all_df['limitUp']]['code'].tolist()
        [dayResult.add_limit_up(c) for c in codes]

        codes = all_df[(all_df['date'] == tradingDay) & all_df['limitDown']]['code'].tolist()
        [dayResult.add_limit_down(c) for c in codes]

        dayLimitUp = self.__calcStrongContinued(all_df, tradingDay)
        dayResult.setDayLimitUp(dayLimitUp)

        amount = IndexCollector().queryShSzAmount(tradingDay)
        dayResult.setShSzAmount(amount)

        dayResult.toString()
        return dayResult

    def __calcStrongContinued(self, all_df, tradingDay):
        print("__calcStrongContinued")
        # calc 當天漲停向前找 strong,continued
        startDay = Objects.getDay(-1 * StockProcessor.CALC_PRE_DAYS)
        a = [d for d in list(filter(lambda day: startDay <= day and day <= tradingDay, self.calcTradingDays))]

        limitUpCodes = all_df[(all_df['date'] == tradingDay)]['code'].tolist()
        # strong_day strong_limit_up
        # loop stock
        for code in limitUpCodes:
            if code[-1] == '1' and code[-2] == '1':
                print(".", end=" ")
            sub_df = (all_df[(all_df['date'] <= tradingDay) & (all_df['code'] == code)]
                      .sort_values(by='date', ascending=False).reset_index())

            all_strong = []
            # loop day
            for index, row in sub_df.iterrows():
                strong = self.findStrongDayStrongLimitUp(sub_df[index:-1])
                all_strong.append(strong)
            max_tuple = self.getMaxStrong(all_strong)
            strong_str = 'd{day}u{up}'.format(day=max_tuple[0], up=max_tuple[1])

            rowIndex = all_df[(all_df['date'] == tradingDay) & (all_df['code'] == code)].index
            all_df.loc[rowIndex, 'strong'] = strong_str
            # print(all_df.loc[rowIndex, 'strong'])
            # print(all_df[(all_df['code'] == code) & (all_df['strong'] >= 'd1u1')])

        return all_df[(all_df['date'] == tradingDay) & (all_df['strong'] >= 'd1u1')].copy()

    ## TODO
    def getMaxStrong(self, all_strong):
        def max_tuple(a, b):
            return a[0] + a[1] > b[0] + b[1]

        sorted_items = sorted(all_strong, key=cmp_to_key(max_tuple))
        return sorted_items[0]

    def findStrongDayStrongLimitUp(self, df=pd.DataFrame()):
        day = 0
        limitUp = 0
        for index, row in df.iterrows():
            # max calc day
            if index == StockProcessor.CALC_PRE_DAYS:
                break
            if index >= 1 and not row['limitUp']:
                break
            day = day + 1
            if row['limitUp']:
                limitUp = limitUp + 1
        # (1,0) ignore
        return (0, 0) if day == 1 and limitUp == 0 else (day, limitUp)

    def isLimitUp(self, row):
        return Objects.isLimit(row['code'], row['preclose'], row['close'], True)

    def isLimitDown(self, row):
        return Objects.isLimit(row['code'], row['preclose'], row['close'], False)

    def isBreakDown(self, row):
        return (Objects.isLimit(row['code'], row['preclose'], row['high'], True) and
                row['high'] != row['close'])


if __name__ == '__main__':
    processor = StockProcessor()
    START_DAY = '2025-04-07'
    END_DAY = '2025-04-08'
    for i in range(100):
        curday = Objects.getDiffDay(START_DAY, i)
        if curday == END_DAY:
            exit()
        processor.process(startDay=curday, endDay=curday)
